#!/usr/bin/env python
# Copyright (C) 2017 vFeed IO
# This file is part of vFeed Correlated Vulnerability & Threat Database Python Wrapper  - https://vfeed.io
# See the file 'LICENSE' for copying permission.


import json

from config.constants import edb_url
from lib.common.database import Database


class CveExploit(object):
    def __init__(self, cve):
        self.cve = cve.upper()
        (self.cur, self.query) = Database(self.cve).db_init()
        self.data = Database(self.cve, self.cur, self.query).check_cve()

    def get_msf(self):
        """ Metasploit method
        :return: JSON response with Msf ID, file and title
        """
        self.msf = []
        self.cur.execute(
            'SELECT * FROM map_cve_msf WHERE cveid=?', self.query)

        for self.data in self.cur.fetchall():
            item = {"id": self.data[0], "file": self.data[1], "title": self.data[2]}
            self.msf.append(item)

        if len(self.msf) != 0:
            return json.dumps(self.msf, indent=2, sort_keys=True)
        else:
            return json.dumps(None)

    def get_saint(self):
        """ Saint  method
        :return: JSON response with Saint ID, file and title
        """
        self.saint = []
        self.cur.execute(
            'SELECT * FROM map_cve_saint WHERE cveid=?', self.query)

        for self.data in self.cur.fetchall():
            item = {"id": self.data[0], "title": self.data[1], "url": self.data[2]}
            self.saint.append(item)

        if len(self.saint) != 0:
            return json.dumps(self.saint, indent=2, sort_keys=True)
        else:
            return json.dumps(None)

    def get_edb(self):
        """ Exploit Database  method
        :return: JSON response with EDB ID, file and url
        """
        self.edb = []
        self.cur.execute(
            'SELECT * FROM map_cve_exploitdb WHERE cveid=?', self.query)

        for self.data in self.cur.fetchall():
            item = {"id": self.data[0], "file": self.data[1], "url": edb_url + str(self.data[0])}
            self.edb.append(item)

        if len(self.edb) != 0:
            return json.dumps(self.edb, indent=2, sort_keys=True)
        else:
            return json.dumps(None)

    def get_d2(self):
        """ D2 elliot  method
        :return: JSON response with D2 Elliot title and file
        """
        self.d2 = []
        self.cur.execute(
            'SELECT * FROM map_cve_d2 WHERE cveid=?', self.query)

        for self.data in self.cur.fetchall():
            item = {"title": self.data[0], "url": self.data[1]}
            self.d2.append(item)

        if len(self.d2) != 0:
            return json.dumps(self.d2, indent=2, sort_keys=True)
        else:
            return json.dumps(None)
